from sklearn.metrics import classification_report
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.callbacks import ModelCheckpoint

from self_model import Self_Net
from dataset_load import load_dataset
from train_pac.cnn.pic_it import pic_it

# from captcha_get import captcha_product


def train_model(classes: int, data_path: str, width: int = 32,
                height: int = 32, depth: int = 3):
    """
    模型训练

    :param data_path: 数据集路径
    :param classes: 数据集包含类别数
    :param depth: 图片通道数
    :param width: 图片宽
    :param height: 图片高
    :return:
    """
    # 加载数据集
    data, label = load_dataset(width, height, data_path)
    data = data / 255.0

    # 编码
    lb = LabelBinarizer()
    label = lb.fit_transform(label)

    # 分割数据集
    x_train, x_test, y_train, y_test = train_test_split(data, label, test_size=0.25)

    # 验证集
    _, v_x_train, _, v_y_train = train_test_split(x_train,
                                                  y_train, test_size=0.25)

    #
    # num = int(len(x_train) * 0.25)
    # v_x_train = x_train[:num]
    # v_y_train = y_train[:num]

    # 保存最佳模型
    check_point = ModelCheckpoint("captcha_best.hdf5", monitor="val_loss",
                                  mode="min", save_best_only=True, verbose=1)
    callbacks = [check_point]

    # 编译模型
    print("[info]开始编译模型...")
    model = Self_Net.build(width=width, height=height, depth=depth,
                           classes=classes)
    sgd = SGD(learning_rate=0.01, decay=0.01 / 40)

    model.compile(loss="categorical_crossentropy", optimizer=sgd,
                  metrics=["acc"])

    # 训练
    print("[info]开始训练...")
    record = model.fit(x_train, y_train, validation_data=(v_x_train, v_y_train),
                       batch_size=64, epochs=40, verbose=1, callbacks=callbacks)

    # 评估网络
    print("[info]开始评估网络")
    predictions = model.predict(x_test, batch_size=64)
    print(classification_report(y_test.argmax(1),
                                predictions.argmax(1),
                                target_names=[str(i) for i in lb.classes_]))
    pic_it(epochs=40, record=record)
# if __name__ == '__main__':
#     train_pac()
